package com.le.edu.listener;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.google.gson.Gson;
import com.le.edu.config.AIProperties;
import com.le.edu.domain.vo.AIVO;
import com.le.edu.utils.WebSocketRemoteContainerUtil;
import lombok.Data;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import org.springframework.data.redis.core.RedisTemplate;

import java.io.IOException;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @author Le
 * @create 2024-04-29 17:39
 */
public class AIWebSocketListener extends WebSocketListener {
    private Long userId;
    private RedisTemplate<String, Object> redisTemplate;
    private ThreadPoolExecutor executor;
    private AIProperties properties;
    private Boolean wsCloseFlag;
    public static final Gson gson = new Gson();

    public static String totalAnswer = ""; // 大模型的答案汇总
    public static String question = ""; // 用户的问题汇总

    public AIWebSocketListener() {
    }

    public AIWebSocketListener(Long userId,
                               RedisTemplate<String, Object> redisTemplate,
                               AIProperties properties,
                               ThreadPoolExecutor executor,
                               Boolean wsCloseFlag,
                               String question) {
        this.userId = userId;
        this.redisTemplate = redisTemplate;
        this.properties = properties;
        this.executor = executor;
        this.wsCloseFlag = wsCloseFlag;
        this.question = question;
    }

    @Override
    public void onOpen(WebSocket webSocket, Response response) {
        super.onOpen(webSocket, response);
        CompletableFuture.runAsync(() -> {
            try {
                JSONObject requestJson = new JSONObject();

                JSONObject header = new JSONObject();  // header参数
                header.put("app_id", properties.getAppid());
                header.put("uid", UUID.randomUUID().toString().substring(0, 10));

                JSONObject parameter = new JSONObject(); // parameter参数
                JSONObject chat = new JSONObject();
                chat.put("domain", "generalv2");
                chat.put("temperature", 0.5);
                chat.put("max_tokens", 4096);
                parameter.put("chat", chat);

                JSONObject payload = new JSONObject(); // payload参数
                JSONObject message = new JSONObject();
                JSONArray text = new JSONArray();

                // 最新问题
                RoleContent roleContent = new RoleContent();
                roleContent.setRole("user");
                roleContent.setContent(question);
                text.add(JSON.toJSON(roleContent));


                message.put("text", text);
                payload.put("message", message);


                requestJson.put("header", header);
                requestJson.put("parameter", parameter);
                requestJson.put("payload", payload);
                webSocket.send(requestJson.toString());
                // 等待服务端返回完毕后关闭
                while (true) {
                    Thread.sleep(200);
                    if (wsCloseFlag) {
                        break;
                    }
                }
                webSocket.close(1000, "");
            } catch (Exception e) {
                e.printStackTrace();
            }
        }, executor);
    }

    @Override
    public void onMessage(WebSocket webSocket, String text) {
        JsonParse myJsonParse = gson.fromJson(text, JsonParse.class);
        if (myJsonParse.header.code != 0) {
            System.out.println("发生错误，错误码为：" + myJsonParse.header.code);
            System.out.println("本次请求的sid为：" + myJsonParse.header.sid);
            webSocket.close(1000, "");
        }
        List<Text> textList = myJsonParse.payload.choices.text;
        for (Text temp : textList) {
            WebSocketRemoteContainerUtil.sendOneMsg(userId, JSON.toJSONString(new AIVO(0, temp.content)));
            totalAnswer = totalAnswer + temp.content;
        }

        if (myJsonParse.header.status == 2) {
            // 可以关闭连接，释放资源
            wsCloseFlag = true;
            WebSocketRemoteContainerUtil.sendOneMsg(userId, JSON.toJSONString(new AIVO(1, "")));
        }
    }

    @Override
    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
        super.onFailure(webSocket, t, response);
        try {
            if (null != response) {
                int code = response.code();
                System.out.println("onFailure code:" + code);
                System.out.println("onFailure body:" + response.body().string());
                if (101 != code) {
                    System.out.println("connection failed");
                    System.exit(0);
                }
            }
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    class JsonParse {
        Header header;
        Payload payload;
    }

    class Header {
        int code;
        int status;
        String sid;
    }

    class Payload {
        Choices choices;
    }

    class Choices {
        List<Text> text;
    }

    class Text {
        String role;
        String content;
    }

    @Data
    class RoleContent {
        String role;
        String content;
    }
}
